package com.binance.actuator.indicator;

import com.binance.actuator.model.MultiDateSourceInfo;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.springframework.boot.actuate.health.Health;
import org.springframework.boot.actuate.health.HealthIndicator;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.util.StopWatch;

import javax.sql.DataSource;
import java.util.List;
import java.util.Map;

// 针对shardingjdbc这种多数据源
public class MultiDataSourceHealthIndicator implements HealthIndicator {

    private static final String DEFAULT_QUERY = "SELECT 1";// 心跳语句

    private Map<String, DataSource> multiDataSource;

    private Map<String, JdbcTemplate> multiJdbcTemplate;

    public MultiDataSourceHealthIndicator(Map<String, DataSource> dataSource) {
        this.multiDataSource = dataSource;
        initJdbcTemplate();
    }

    /**
     * 初始化jdbctemplate
     */
    private void initJdbcTemplate() {
        if (MapUtils.isEmpty(multiDataSource)) {
            return;
        }
        multiJdbcTemplate = Maps.newHashMap();
        for (Map.Entry<String, DataSource> entry : multiDataSource.entrySet()) {
            JdbcTemplate jdbcTemplate = new JdbcTemplate();
            jdbcTemplate.setDataSource(entry.getValue());
            multiJdbcTemplate.put(entry.getKey(), jdbcTemplate);
        }
    }


    @Override
    public Health health() {
        List<MultiDateSourceInfo> multiDateSourceInfoList = checkDateSource();
        List<MultiDateSourceInfo> errorDataSourceInfoList = Lists.newArrayList();
        Health.Builder builder = new Health.Builder();
        builder.withDetail("multiDataSourceInfo", multiDateSourceInfoList);
        for (MultiDateSourceInfo multiDateSourceInfo : multiDateSourceInfoList) {
            if ("DOWN".equalsIgnoreCase(multiDateSourceInfo.getStatus())) {
                errorDataSourceInfoList.add(multiDateSourceInfo);
            }
        }
        if(CollectionUtils.isNotEmpty(errorDataSourceInfoList)){
            builder.withDetail("errorMessage", errorDataSourceInfoList);
            return builder.down().build();
        }
        return builder.up().build();
    }

    private String queryDB(JdbcTemplate jdbcTemplate) {
        String result = jdbcTemplate.queryForObject(DEFAULT_QUERY, String.class);
        return result;
    }

    private List<MultiDateSourceInfo> checkDateSource() {
        List<MultiDateSourceInfo> multiDateSourceInfoList = Lists.newArrayList();
        for (Map.Entry<String, JdbcTemplate> entry : multiJdbcTemplate.entrySet()) {
            StopWatch stopWatch = new StopWatch();
            try {
                stopWatch.start();
                String result = queryDB(entry.getValue());
                stopWatch.stop();
                MultiDateSourceInfo multiDateSourceInfo = new MultiDateSourceInfo();
                multiDateSourceInfo.setDbName(entry.getKey());
                multiDateSourceInfo.setStatus("UP");
                multiDateSourceInfo.setQueryResult(result);
                multiDateSourceInfo.setTotalTimeSecond(String.valueOf(stopWatch.getTotalTimeSeconds()));
                multiDateSourceInfoList.add(multiDateSourceInfo);
            } catch (Exception e) {
                if (stopWatch.isRunning()) {
                    stopWatch.stop();
                }
                MultiDateSourceInfo exceptionDataSourceInfo = new MultiDateSourceInfo();
                exceptionDataSourceInfo.setDbName(entry.getKey());
                exceptionDataSourceInfo.setStatus("DOWN");
                exceptionDataSourceInfo.setErrorMessage(e.toString());
                exceptionDataSourceInfo.setTotalTimeSecond(String.valueOf(stopWatch.getTotalTimeSeconds()));
                multiDateSourceInfoList.add(exceptionDataSourceInfo);
            }
        }
        return multiDateSourceInfoList;
    }
}
